# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: chunhui

import collections
import datetime
import pandas
import numpy

from absl.flags import FLAGS
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2
from coin.strategy.mm.tool.pnl_plot import launch_all as dump_log_job
from coin.strategy.mm.tool.signal_plot_base import dump_funding_rate
import coin.pnl.pnl_compact_stat as pcstat
import coin.pnl.sim_stat_calc2 as sscalc2
import coin.strategy.mm.tool.archive_base as abase


def gen_symbol_product_map(orddf):
  mapping = {}
  for symbol, sym_orddf in orddf.groupby('symbol'):
    first_row = sym_orddf.iloc[0]
    product = generate_product_from_str2(
        first_row['market_type'],
        first_row['exchange_type'],
        None,
        first_row['symbol'],
        pandas.DatetimeIndex([first_row['timestamp']]).to_pydatetime()[0])
    relative_norm = product.subscription_symbol.replace('CURRENT_', '')
    mapping[symbol] = product
    mapping[relative_norm] = product
  return mapping


def consolidate_bbo_dicts(bbo_dicts, symbol_product_map):
  new_dict = collections.defaultdict(list)
  for symbol, bbodfs in bbo_dicts.items():
    if symbol in symbol_product_map:
      relative_norm = symbol_product_map[symbol].subscription_symbol.replace('CURRENT_', '')
      new_dict[relative_norm] += bbodfs
    else:
      new_dict[symbol] += bbodfs
  return new_dict


def read_live_log():
  baseflags = abase.get_base_flags()
  fsflags = abase.get_feed_subscription_flags()
  osflags = abase.get_order_subscription_flags()

  args = tuple([baseflags, fsflags, osflags])
  rettuples = []
  rettuples.append(dump_log_job(*args))

  bbodf_dicts = collections.defaultdict(list)
  orddfs = []
  for orddf, bbodf_dict in rettuples:
    if orddf is None:
      continue
    orddfs.append(orddf)
    for (_, symbol), bbodf in bbodf_dict.items():
      bbodf_dicts[symbol].append(bbodf)
  orddf = pandas.concat(orddfs, axis=0, sort=False).reset_index(drop=True)
  orddf = orddf.sort_values('timestamp').reset_index(drop=True)

  sym_product_map = gen_symbol_product_map(orddf)
  orddf['symbol'] = orddf['symbol'].map(
      lambda x: sym_product_map[x].subscription_symbol.replace('CURRENT_', ''))
  bbodf_dicts = consolidate_bbo_dicts(bbodf_dicts, sym_product_map)

  bbodf_dict = {}
  for symbol, bbodfs in bbodf_dicts.items():
    bbodf = pandas.concat(bbodfs, axis=0, sort=False).reset_index(drop=True)
    if len(bbodf) > 0:
      bbodf = bbodf.sort_values('timestamp').reset_index(drop=True)
      bbodf = bbodf[numpy.log(bbodf['bid0_price'] / bbodf['ask0_price']).abs() < 0.05]
    bbodf_dict[symbol] = bbodf

  sstats_dict = {}
  timedrets = []

  first_event_ts = None
  trading_dates = abase.get_trading_dates(baseflags.trading_date)
  for symbol, sym_orddf in orddf.groupby('symbol'):
    if FLAGS.symbol is not None and symbol.find(FLAGS.symbol) < 0:
      continue
    sym_orddf = sym_orddf[:-1].copy()  # Remove last EOD virtual fill
    event_ts = sym_orddf[sym_orddf.type == 'ORDER_FILLED'].iloc[0]['timestamp']
    if first_event_ts is None or event_ts < first_event_ts:
      first_event_ts = event_ts
    product = sym_product_map[symbol]
    if (sym_bbodf := bbodf_dict.get(symbol, None)) is None:
      print(f"{symbol} is ignored")
      continue
    sym_orddf['fill_type'] = sym_orddf['fill_type'].str.replace('_FILL_TYPE', '')

    funding_df = dump_funding_rate(product, trading_dates[0],
                                   trading_dates[-1] + datetime.timedelta(days=1),
                                   FLAGS.bbo_cache_root)

    compact_stat = pcstat.PnlCompactStat.from_orderdf(
        product=product,
        orderdf=sym_orddf
    )
    sresp = sscalc2.SimResponseCompact(
        param_label=None,
        compact_stat=compact_stat,
        miscinfo={'initial_pos': 0, "funding_df": funding_df},
        bbodf=sym_bbodf,
        product=product)
    simstat = sscalc2.get_stat_df(sresp, {})
    sstats_dict[symbol] = [simstat]
    if (timedret := simstat.get('timedret', None)) is not None:
      timedrets.append(timedret)
  return sstats_dict, int(first_event_ts)
